import numpy as np
from flask import Blueprint, request, jsonify

from src import faiss_index
from src.routes.model import IndexRequest

idx_bp = Blueprint('idx', __name__)


def post_construct(name: str, dim: int, index_factory: str, hdfs_path: str = None, local_path: str = None):
    faiss_index.post_construct(name, dim, index_factory, hdfs_path, local_path)


@idx_bp.post('/index')
def index():
    idx: IndexRequest = request.get_json()
    index_name = idx.get('index')
    if faiss_index.get_index(index_name) is not None:
        return {'success': False, 'msg': 'index already exists'}
    faiss_index.set_index(index_name, faiss_index.INDEX_PLACEHOLDER)
    faiss_index.async_index_register(idx.get('index'), idx.get('dim'), idx.get('index_factory'),
                                     idx.get('hdfs_path', None), idx.get('local_path', None))
    return {'success': True, 'msg': 'index building,please check /info/<idx_name>'}


@idx_bp.get('/infos')
def infos():
    return {'success': True, 'data': faiss_index.infos()}


@idx_bp.get('/info/<index_name>')
def info(index_name=None):
    # 查看目标index的索引情况：存在、量级、索引构建状态、索引构建参数
    if index_name is None:
        return infos()
    index_instance = faiss_index.get_index(index_name)
    if index_instance == faiss_index.INDEX_PLACEHOLDER:
        return {'success': True, 'msg': 'index initializing'}
    elif index_instance is None:
        return {'success': True, 'msg': 'not found index'}
    return {'success': True, 'data': index_instance.info()}


@idx_bp.post('/search')
def search():
    data = request.get_json()
    # 校验 index、query、topk
    index_name = data.get('index')
    query = data.get('query')
    list_query = np.array([float(x) for x in query.split(",")]).reshape(1, -1)
    topk = data.get('topk')
    idx = faiss_index.get_index(index_name)
    if idx == faiss_index.INDEX_PLACEHOLDER:
        return {'success': False, 'msg': 'index initializing'}
    elif idx is None:
        return {'success': False, 'msg': 'not found index'}
    # 查询
    res = idx.search(list_query, topk)
    return {'success': True, 'data': [e.to_dict() for e in res]}
